# -*- coding: utf-8 -*-

import os
from PIL import Image
import shutil

def check_charset(file_path):
    """
    自动检测文件编码，避免编码错误
    """
    import chardet
    with open(file_path, "rb") as f:
        data = f.read(4)
        charset = chardet.detect(data)['encoding']
    return charset


def convert(size, xmin, ymin, xmax, ymax):
    """
    VOC格式边界框转YOLO格式
    size: (width, height)
    xmin,ymin,xmax,ymax: 目标框坐标
    返回: (x_center, y_center, width, height)均归一化到[0,1]
    """
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (xmin + xmax) / 2.0
    y = (ymin + ymax) / 2.0
    w = xmax - xmin
    h = ymax - ymin
    return (x * dw, y * dh, w * dw, h * dh)


def extract_labels_images(outpath_txt, outpath_jpg, ori_data_path, origin_txt_path):
    """
    读取WiderPerson自定义标注txt，转换成YOLO标签格式，
    同时移动图片到新目录
    """
    with open(origin_txt_path, 'r') as f:
        img_ids = [x.strip() for x in f.readlines()]

    for img_id in img_ids:
        img_path = os.path.join(ori_data_path, 'Images', img_id + '.jpg')

        # 读取图片尺寸
        with Image.open(img_path) as img:
            img_size = img.size  # (width, height)

        ans = ''
        label_path = os.path.join(ori_data_path, 'Annotations', img_id + '.jpg.txt')
        outpath = os.path.join(outpath_txt, img_id + '.txt')

        # 读取标注txt
        with open(label_path, encoding=check_charset(label_path)) as file:
            count_line = file.readline()
            count = int(count_line.strip())  # 行人数
            line = file.readline()
            while line:
                parts = line.strip().split(' ')
                if len(parts) < 5:
                    line = file.readline()
                    continue
                cls = int(parts[0])
                if cls == 1:  # 只保留类别1，行人
                    xmin = float(parts[1])
                    ymin = float(parts[2])
                    xmax = float(parts[3])
                    ymax = float(parts[4])
                    bb = convert(img_size, xmin, ymin, xmax, ymax)
                    ans += '1 ' + ' '.join(f'{a:.6f}' for a in bb) + '\n'
                line = file.readline()

        # 写YOLO标签文件
        with open(outpath, 'w') as f:
            f.write(ans)

        # 移动图片到新目录（如果想保留原图，改为shutil.copy）
        shutil.move(img_path, os.path.join(outpath_jpg, img_id + '.jpg'))


def write_label(otxt_path, ntxt_path):
    """
    把标签txt中类别统一转为0（YOLO类别从0开始）
    """
    for file_name in os.listdir(otxt_path):
        otxt = os.path.join(otxt_path, file_name)
        ntxt = os.path.join(ntxt_path, file_name)

        with open(otxt, 'r', encoding='utf-8') as f:
            lines = f.readlines()

        new_lines = []
        for line in lines:
            if line.strip() == '':
                continue
            parts = line.strip().split(' ')
            # 把所有类别改为0
            new_line = '0 ' + ' '.join(parts[1:]) + '\n'
            new_lines.append(new_line)

        with open(ntxt, 'w') as f:
            f.writelines(new_lines)


def write_train_val_txt(labels_path, txt_path, image_set):
    """
    生成train.txt或val.txt文件，内容是图片相对路径列表
    """
    image_ids = []
    label_dir = os.path.join(labels_path, image_set)
    for file_name in os.listdir(label_dir):
        image_id = os.path.splitext(file_name)[0]
        image_ids.append(image_id)

    list_file_path = os.path.join(txt_path, image_set + '.txt')
    with open(list_file_path, 'w') as f:
        for img_id in image_ids:
            f.write(f'./images/{image_set}/{img_id}.jpg\n')


if __name__ == '__main__':
    print('WiderPerson数据集yolo格式文件抽取程序启动：')

    # 原始数据集路径（请根据实际修改）
    ori_data_path = './WiderPerson'
    # 输出数据路径
    out_data_path = './WiderPerson/WiderPerson_yolo'

    sets = ['train', 'val']

    for data_set in sets:
        print(f'（1）{data_set}数据类别抽取中......')
        outpath_txt = os.path.join(out_data_path, 'label', data_set)
        outpath_jpg = os.path.join(out_data_path, 'images', data_set)
        os.makedirs(outpath_txt, exist_ok=True)
        os.makedirs(outpath_jpg, exist_ok=True)
        origin_txt_path = os.path.join(ori_data_path, f'{data_set}.txt')

        extract_labels_images(outpath_txt, outpath_jpg, ori_data_path, origin_txt_path)

        print(f'（2）{data_set}_label文件写入中......')
        otxt_path = os.path.join(out_data_path, 'label', data_set)
        ntxt_path = os.path.join(out_data_path, 'labels', data_set)
        os.makedirs(ntxt_path, exist_ok=True)
        write_label(otxt_path, ntxt_path)

        print(f'（3）{data_set}.txt写入中......')
        labels_path = os.path.join(out_data_path, 'labels')
        txt_path = os.path.join(out_data_path, 'labels')
        write_train_val_txt(labels_path, txt_path, data_set)

    print('（4）删除多余文件中......')
    shutil.rmtree(os.path.join(out_data_path, 'label'))
    print('数据集抽取完成！！！')
